package mllib;

import java.util.ArrayList;
import org.apache.commons.math.linear.*;

public class KMeansEstimator implements Estimator
{
    private int k;
    private ArrayList<RealVector> means;


    public KMeansEstimator(TrainingSet ts, int k)
    {
        this.k = k;
        means = new ArrayList<RealVector>();
        calcMeans(ts);
       
        calcMeans(ts);
    }

    private void calcMeans(TrainingSet ts)
    {
        //select random centroid points
        for( int i = 0; i < k; i++)
            means.add( ts.getSample( (int)(Math.random()*ts.dimension() )) );

        int[] nearestNeighbor = new int[ts.dimension()];
        boolean changed;

        do{     
            changed = false;   

            //fill in nearest neighbors
            for(int t = 0; t < ts.dimension(); t++)
            {
                RealVector x = ts.getSample(t);
                nearestNeighbor[t] = 0;
                for(int i = 1; i < means.size(); i++)
                {
                     if( x.getDistance(means.get(i)) < x.getDistance( means.get(nearestNeighbor[t])))
                        nearestNeighbor[t] = i;
                }
            }

            //calculate average of points that have m as nearest mean
            for(int idx = 0; idx < means.size(); idx++)
            {   
                int count = 0;
                RealVector newM = new ArrayRealVector(ts.dimension());
                for( int x = 0; x < ts.dimension(); x++)
                {
                    if( nearestNeighbor[x] == idx)
                    {
                        count++;
                        newM = newM.add( ts.getSample(x));
                    }
                }
                if( count == 0) 
                    means.remove( idx);
                else 
                    newM = newM.mapDivideToSelf( count);

                changed = changed || sameVector(newM, means.get(idx));
                means.add(idx, newM);        
            }

        }while( changed );
            
    }

    private boolean sameVector(RealVector a, RealVector b)
    {
        for( int i = 0; i < a.getDimension(); i++)
            if( a.getEntry(i) != b.getEntry(i) )
                return false;
        return true;
    }

    public void toFile( String filename)
    {
        //TODO 
    }

    public double prob( RealVector x) throws Exception
    {
        double dist = 0;
        for( RealVector m: means)
            dist += m.getDistance( x);
        return dist;
    }


}
